# coding: utf-8
import os
import json
import cv2
import numpy as np
from tqdm import tqdm

def extract_defects_from_dataset(
    data_root, output_root, target_classes
):
    """
    从指定类别文件夹中提取缺陷区域和掩码，保存到 output_root。
    每个 JSON 仅包含一个 shape, 多个 shape 的跳过。
    """
    for cls in target_classes:
        class_dir = os.path.join(data_root, cls)
        if not os.path.exists(class_dir):
            print(f"[警告] 类别目录不存在: {class_dir}")
            continue

        output_class_dir = os.path.join(output_root, cls)
        os.makedirs(output_class_dir, exist_ok=True)

        files = os.listdir(class_dir)
        json_files = [f for f in files if f.endswith(".json")]

        for json_file in tqdm(json_files, desc=f"处理类别: {cls}"):
            json_path = os.path.join(class_dir, json_file)
            image_name = json_file.replace(".json", ".jpg")
            image_path = os.path.join(class_dir, image_name)

            if not os.path.exists(image_path):
                print(f"❌ 图像缺失：{image_path}")
                continue

            with open(json_path, 'r') as f:
                try:
                    data = json.load(f)
                except:
                    print(f"⚠️ 无法解析 JSON: {json_path}")
                    continue

            shapes = data.get("shapes", [])
            if len(shapes) != 1:
                # 多个 shape 的直接跳过
                continue

            shape = shapes[0]
            points = np.array(shape['points'], dtype=np.int32)

            image = cv2.imread(image_path)
            if image is None:
                print(f"❌ 图像读取失败：{image_path}")
                continue

            # 创建掩码并裁剪
            mask = np.zeros(image.shape[:2], dtype=np.uint8)
            cv2.fillPoly(mask, [points], 255)

            x, y, w, h = cv2.boundingRect(points)
            cropped_patch = image[y:y+h, x:x+w]
            cropped_mask = mask[y:y+h, x:x+w]

            base_name = os.path.splitext(json_file)[0]
            patch_path = os.path.join(output_class_dir, f"{base_name}_patch.png")
            mask_path = os.path.join(output_class_dir, f"{base_name}_mask.png")

            cv2.imwrite(patch_path, cropped_patch)
            cv2.imwrite(mask_path, cropped_mask)

        print(f"✅ 类别 {cls} 处理完成，保存于 {output_class_dir}")


if __name__ == '__main__':
    data_root = r"/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10_SD/20250530-SD-small-defect"
    output_root = r"/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10_SD/20250530-SD-small-defect-patch"
    target_classes = ["TG05", "TS28", "TS15"]
    extract_defects_from_dataset(data_root,output_root,target_classes)
